Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.71%. Comparing base (e71d0b4) to head (01f47ff).
⚠️ Report is 11 commits behind head on main.
✅ All tests successful. No failed tests found.

@@            Coverage Diff             @@
##             main    #1722      +/-   ##
==========================================
- Coverage   88.51%   84.71%   -3.81%     
==========================================
  Files         137      137              
  Lines       11527    11518       -9     
==========================================
- Hits        10203     9757     -446     
- Misses       1324     1761     +437     

Thanks for the review @dgedon

To clarify: the PR description refers to an earlier change (PR #1370) that converted proposal from torch.Distribution to Callable. This PR is about making that Callable type more precise with Protocols.

On whether Protocols are necessary vs Callable[[torch.Size], Tensor]: the key issue is that SampleProposal needs to accept **kwargs. Looking at actual usage:

  candidates = proposal(                                                                                                                                             
      torch.Size((sampling_batch_size,)),                                                                                                                            
      **proposal_sampling_kwargs,  # passes {"condition": x}, {"predictor": ..., "corrector": ...}, etc.                                                             
  )

The proposals (posterior_estimator.sample, _sample_via_diffusion, etc.) all require additional kwargs like condition. You can't express (sample_shape: torch.Size, **kwargs) -> Tensor with a simple Callable type hint — that's the main reason for the Protocol.

For AcceptRejectFn, you're right that Callable[[Tensor], Tensor] would technically suffice since it's only ever called with theta. I kept it as a Protocol for consistency and more explicit documentation, but happy to simplify if you prefer.